# %% Vision-based formation control of quadrotors
#
# Notes:
# See the tutorial for explanations.
# The simulation environment in UE4 must be running before executing this code.
#
# Jing Tian,2025.
# 原始作者
# (C) Kaveh Fathian, 2017-2018.（原始作者）
#
# Website: https://sites.google.com/view/kavehfathian


# %%

import Include.multirotor.setup_path
import airsim

import os
import time
import numpy as np
import pprint
import time
import math
import matlab

import matlab.engine

matlabEngine = matlab.engine.start_matlab()  # Connect to matlab API

# %% Desired formation parameters

# Total number of UAVs
numUAV = 3

# Desired formation (equilateral triangle)
qs = matlab.double([[0, 4, 2], [0, 0, 3]])

# Adjacency matrix
Adjm = matlab.double([[0, 1, 1], [1, 0, 1], [1, 1, 0]])

# Initial positions of the quads
# NOTE: THIS MUST MATCH WITH THE SETTING FILE!
pos0 = np.zeros((numUAV, 3))
for i in range(numUAV):
    pos0[i, 0] = 588.0
    pos0[i, 1] = -28.0 + 4.0 * i
    pos0[i, 2] = -125.0


# %% Find formaiton control gains

# Formation control gains
Am = matlabEngine.FindGains_Ver1_4(qs, Adjm, nargout=1)
# print(Am)

# Conver Am to numpy array
A = np.asarray(Am)
Adj = np.asarray(Adjm)
print("Gain matrix calculated.")


# %% Connect to the AirSim simulator

client = airsim.MultirotorClient()
client.confirmConnection()

for i in range(numUAV):
    name = "UAV" + str(i + 1)
    client.enableApiControl(True, name)
    client.armDisarm(True, name)
print("All UAVs have been initialized.")


# Hover
time.sleep(2)

tout = 3  # Timeout in seconds
spd = 2  # Speed

print("taking off...")
for i in range(numUAV):
    name = "UAV" + str(i + 1)
    print("Hovering", name)
    client.hoverAsync(vehicle_name=name)
    client.moveToPositionAsync(0, 0, -1, spd, timeout_sec=tout, vehicle_name=name)
print("All UAVs are hovering.")


# %% Increase altitude

tout = 10.0  # Timeout in seconds
spd = 4.0  # Speed
alt = -20.0  # Altitude

time.sleep(0.5)

for i in range(numUAV):
    name = "UAV" + str(i + 1)
    print("Moving", name)
    client.moveToPositionAsync(0, 0, alt, spd, timeout_sec=tout, vehicle_name=name)
print("UAVs reached desired altitude")


# %% Formation control loop

dcoll = 1.5  # Collision avaoidance activation distance
rcoll = 0.7  # Collision avaoidance circle radius
gain = 1.0 / 3  # Control gain
alt = -20.0  # UAV altitude
duration = 0.5  # Max duration for applying input
vmax = 0.1  # Saturation velocity
save = 0  # Set to 1 to save onboard images, otherwise set to 0


# Initial Pause time
time.sleep(0.5)

# Get Image data (Initialization)
for i in range(numUAV):
    name = "UAV" + str(i + 1)
    imgs = client.simGetImages(
        [airsim.ImageRequest("1", airsim.ImageType.Scene, False, False)],
        vehicle_name=name,
    )  # Scene vision image in uncompressed RGBA array
    img = imgs[0]
    imgWidth = img.width
    imgHeight = img.height
    imgData = img.image_data_uint8

    # 判断图像格式并转换为 RGBA
    expected_length = imgWidth * imgHeight * 4
    actual_length = len(imgData)
    if actual_length == expected_length:
        print(f"{name} 图像格式可能是 RGBA，无需转换。")
    elif actual_length == imgWidth * imgHeight * 3:
        print(f"{name} 图像格式可能是 RGB，转换为 RGBA...")
        rgba_img = []
        for j in range(0, actual_length, 3):
            r, g, b = imgData[j], imgData[j + 1], imgData[j + 2]
            rgba_img.extend([r, g, b, 255])  # 添加完全不透明的 alpha 通道
        imgData = bytes(rgba_img)
    elif actual_length == imgWidth * imgHeight:
        print(f"{name} 图像格式可能是灰度图，转换为 RGBA...")
        rgba_img = []
        for gray in imgData:
            rgba_img.extend([gray, gray, gray, 255])  # 复制灰度值到 RGB 通道并添加不透明 alpha 通道
        imgData = bytes(rgba_img)
    else:
        print(f"{name} 图像格式未知，无法转换。")

    if i == 0:
        imgArray = imgData
    else:
        imgArray = imgArray + imgData

imgWidth = img.width
imgHeight = img.height

# 判断图像格式（再次确认，转换后应该是 RGBA）
expected_length = numUAV * imgWidth * imgHeight * 4  # RGBA 格式下预期的长度
if len(imgArray) == expected_length:
    print("最终图像格式为 RGBA。")
else:
    print("最终图像格式仍存在问题。")

# 打开一个绘图窗口，绘制轨迹图，以下部分调用matlab
matlabEngine.doublebuffer_win(nargout=0)  # 注意不能直接调用matlab中的set()函数

# Formation control
itr = 0
# e=10000#误差

while True:

    # 获取无人机实时位置，绘制轨迹图
    x1 = client.simGetGroundTruthKinematics(vehicle_name="UAV1").position.x_val
    y1 = client.simGetGroundTruthKinematics(vehicle_name="UAV1").position.y_val
    z1 = client.simGetGroundTruthKinematics(vehicle_name="UAV1").position.z_val
    x2 = client.simGetGroundTruthKinematics(vehicle_name="UAV2").position.x_val
    y2 = client.simGetGroundTruthKinematics(vehicle_name="UAV2").position.y_val
    z2 = client.simGetGroundTruthKinematics(vehicle_name="UAV2").position.z_val
    x3 = client.simGetGroundTruthKinematics(vehicle_name="UAV3").position.x_val
    y3 = client.simGetGroundTruthKinematics(vehicle_name="UAV3").position.y_val
    z3 = client.simGetGroundTruthKinematics(vehicle_name="UAV3").position.z_val
    # 画出三架无人机实时的位置
    matlabEngine.draw_trajectory(
        x1, y1, z1, x2, y2, z2, x3, y3, z3, nargout=0
    )  

    # l1=math.sqrt(math.pow(x1-x2, 2)+math.pow(y1-y2, 2)+math.pow(z1-z2, 2))#三角形边长
    # l2=math.sqrt(math.pow(x2-x1, 2)+math.pow(y2-y1, 2)+math.pow(z2-z1, 2))
    # l3=math.sqrt(math.pow(x3-x1, 2)+math.pow(y3-y1, 2)+math.pow(z3-z1, 2))
    # u_val=(l1+l2+l3)/3#三条边的均值
    # e=(math.pow(l1-u_val,2)+math.pow(l2-u_val,2)+math.pow(l3-u_val,2))/3#计算误差(方差)

    itr = itr + 1
    print("itr = ", itr)
    # print("e = ", e)#输出误差

    if itr == 1:
        print("Starting a prallel pool...")

    # Get UAV positions for collision avoidance
    q = np.zeros(3 * numUAV)  # Preallocate state vectors
    qo = np.zeros(4 * numUAV)  # Preallocate orientation vectors
    qxy = np.zeros(2 * numUAV)
    for i in range(numUAV):
        name = "UAV" + str(i + 1)
        # Get x-y-z coordinates
        pos = client.simGetGroundTruthKinematics(vehicle_name=name)
        qi = np.array([pos.position.x_val, pos.position.y_val, pos.position.z_val])
        qoi = np.array(
            [
                pos.orientation.w_val,
                pos.orientation.x_val,
                pos.orientation.y_val,
                pos.orientation.z_val,
            ]
        )

        # Add initial coordinates
        qd = qi + pos0[i, :]

        # 3D and 2D state vector
        q[3 * i : 3 * i + 3] = qd.copy()
        qxy[2 * i : 2 * i + 2] = np.array([qd[0], qd[1]])
        qo[4 * i : 4 * i + 4] = qoi.copy()

    # Estimate relative positions using onboard images
    for i in range(numUAV):
        name = "UAV" + str(i + 1)
        imgs = client.simGetImages(
            [airsim.ImageRequest("0", airsim.ImageType.Scene, compress=False)],
            vehicle_name=name,
        )  # Scene vision image in uncompressed RGBA array
        img = imgs[0]
        print(
            f"UAV {i+1} image width: {img.width}, height: {img.height}, length: {len(img.image_data_uint8)}"
        )

        imgWidth = img.width
        imgHeight = img.height
        imgData = img.image_data_uint8

        # 判断图像格式并转换为 RGBA
        expected_length = imgWidth * imgHeight * 4
        actual_length = len(imgData)
        if actual_length == expected_length:
            print(f"{name} 图像格式可能是 RGBA，无需转换。")
        elif actual_length == imgWidth * imgHeight * 3:
            print(f"{name} 图像格式可能是 RGB，转换为 RGBA...")
            rgba_img = []
            for j in range(0, actual_length, 3):
                r, g, b = imgData[j], imgData[j + 1], imgData[j + 2]
                rgba_img.extend([r, g, b, 255])  # 添加完全不透明的 alpha 通道
            imgData = bytes(rgba_img)
        elif actual_length == imgWidth * imgHeight:
            print(f"{name} 图像格式可能是灰度图，转换为 RGBA...")
            rgba_img = []
            for gray in imgData:
                rgba_img.extend([gray, gray, gray, 255])  # 复制灰度值到 RGB 通道并添加不透明 alpha 通道
            imgData = bytes(rgba_img)
        else:
            print(f"{name} 图像格式未知，无法转换。")

        if i == 0:
            imgArray = imgData
        else:
            imgArray = imgArray + imgData

    # meng.edit('GetRelativePose_Ver1_7')
    # print(f"imgArray 的长度为: {len(imgArray)}") # 这里打印的长度是有值的，类型为bytes
    print(f"Image width: {imgWidth}, Image height: {imgHeight}")
    print(f"imgArray 的长度为: {len(imgArray)}")  
    imgArray_matlab = matlab.uint8(list(imgArray))
    Qm, Tm, flagm = matlabEngine.GetRelativePose_Ver1_7(
        imgArray_matlab, imgWidth, imgHeight, save, Adjm, itr, nargout=3
    )

    T = np.asarray(Tm)
    flag = np.asarray(flagm).flatten()
    # print(T)

    # Transform recovered coordinates to world frame in order to apply the control.
    # AirSim uses NED (North-East-Down) frame, and the front camera has EDN
    # (East-Down-Notrth) frame. So we need to swao columns of T.
    Tw = np.array([T[:, 2], T[:, 0], T[:, 1]]).T

    # Calculate distributed control
    dqxy = np.zeros(2 * numUAV)  # Preallocate vectors
    for i in range(numUAV):
        if flag[i] == 1:
            # 3D and 2D state vector
            qi = Tw[i * numUAV : (i + 1) * numUAV, :].flatten()
            qxyi = Tw[i * numUAV : (i + 1) * numUAV, 0:2].flatten()

            # Control
            dqxyi = A[i * 2 : i * 2 + 2, :].dot(qxyi)
            dqxy[2 * i : 2 * i + 2] = gain * dqxyi
    if save == 1:
        np.save("SavedData/u" + str(itr), dqxy)  # Save control

    # Collision avoidance  # meng.edit('ColAvoid_Ver2_1')
    um = matlabEngine.ColAvoid_Ver2_1(dqxy.tolist(), qxy.tolist(), dcoll, rcoll, nargout=1)
    u = np.asarray(um).flatten()
    # u = dqxy

    # Saturate velociy control command
    for i in range(numUAV):
        # Find norm of control vector for each UAV
        ui = u[2 * i : 2 * i + 2]
        vel = np.linalg.norm(ui)
        if vel > vmax:
            u[2 * i : 2 * i + 2] = (vmax / vel) * u[2 * i : 2 * i + 2]
    if save == 1:
        np.save("SavedData/um" + str(itr), u)  # Save modified control

    #     Apply control command
    for i in range(numUAV):
        name = "UAV" + str(i + 1)
        client.moveByVelocityZAsync(
            u[2 * i], u[2 * i + 1], alt, duration, vehicle_name=name
        )  # Motion at fixed altitude

    print()

# %% Terminate

matlabEngine.quit()
client.reset()
